# -*- coding: utf-8 -*
# pylint: disable=unsubscriptable-object
import pickle
import matplotlib.pyplot as plt
from pathlib import Path
import numpy as np

from PIL import Image


def read(base_dir):
    base_dir = Path(base_dir)
    images = {1: [], 2: [], 3: []}
    for sub_dir in base_dir.iterdir():
        for jpg in sub_dir.iterdir():
            ix = int(sub_dir.name)
            img = Image.open(jpg)
            images[ix].append([jpg.name, img])
    return images


def scale(images, scale_to):
    for key in images.keys():
        for i, (name, img) in enumerate(images[key]):
            newImg = img.resize((scale_to, scale_to))
            images[key][i][1] = newImg


def toGray(images):
    for key in images.keys():
        for i, (name, img) in enumerate(images[key]):
            newImg = img.convert('L')
            images[key][i][1] = newImg
            print(key, name)


def save(images, to_dir):
    path = Path(to_dir)
    if not path.exists():
        path.mkdir()
    for key in images.keys():
        subdir_path = path / str(key)
        if not subdir_path.exists():
            subdir_path.mkdir()
        for i, (name, img) in enumerate(images[key]):
            jpg_path = subdir_path / name
            print(jpg_path, 'Done!')
            img.save(jpg_path)


def main():
    base_dir = r'./dataset/Origin/题目一'
    to_dir = r'./dataset/Scaled'
    images = read(base_dir)
    scale_to = 900
    scale(images, scale_to)
    # toGray(images)
    save(images, to_dir)


if __name__ == '__main__':
    main()
